/* _NSSOLVER_C_
 * 
 * SOLVES THREE DIMENSIONAL, INCOMPRESSIBLE NAVIER-STOKES
 * EQUATION WITH RECTANGULAR DOMAIN USING FINITE DIFFERENCING 
 * ON A STAGGERED GRID
 * UTILIZES THE PROJECTION METHOD.  THE CONSERVATIVE FORM
 * CONVECTION IS DISCRETIZED WITHOUT THE USE OF UPWIND TREATMENT; 
 * BOTH THE ADVECTION AND DIFUSSION TERMS ARE TREATED EXPLICITLY. 
 * THE CURRENT SETUP IS FOR A LID DRIVEN CAVITY PROBLEM WITH A WELL
 * KNOWN VELOCITY AT THE TOP AND NO-SLIP BOUNDARY CONDIITIONS AT
 * ALL OTHER BOUNDARIES
 * 
 * ORIGINALLY WRITTEN BY DR. ZHI-GANG FENG: UNIVERSITY OF NORTH TEXAS
 * 
 * PORTED TO C AND 3D BY JOHN LEIDEL: GRADUATE STUDENT
 * 
 */

#include "nsproj.h"
#ifdef _OPENMP_
#include <omp.h>
#include <stdlib.h>
#endif

#ifdef _GCC_
void 
set_fpu (unsigned int mode)
{
  asm ("fldcw %0" : : "m" (*&mode));
}
#endif


//----------------------------------------PRESSURE_CORRECTION
// PRESSURE_CORRECTION
// STEP #3 IN PROJECTION METHOD
// CORRECTS PRESSURE FOR ALL VELOCITY VALUES
// PRE: ALL VALUES ARE VALID
// POST: PRESSURE VALUES ARE UPDATED
// RETURNS 0 ON SUCCESS, NONZERO OTHERWISE
//-----------------------------------------------------------
int pressure_correction( double dt, 
			 double dx, 
			 double dy, 
			 double dz, 
			 NSMESH *mesh )
{
    /* VARS */
    int i   = 0;
    int j   = 0;
    int k   = 0;
    /* ---- */

    /* CORRECT U */
#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,j,k) shared(mesh,dt,dx,dy,dz)
#endif
    for( i=1; i<mesh->nx-2; i++ )
    {
	for( j=1; j<mesh->ny; j++ )
	{
	    for( k=1;k<mesh->nz; k++ )
	    {
		mesh->U[i][j][k] = mesh->U[i][j][k] - dt*((mesh->P[i][j-1][k-1]-mesh->P[i-1][j-1][k-1])/dx);
	    }
	}
    }

    /* CORRECT V */
#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,j,k) shared(mesh,dt,dx,dy,dz)
#endif
    for( i=1; i<mesh->nx; i++ )
    {
	for( j=1; j<mesh->ny-2; j++ )
	{
	    for( k=1;k<mesh->nz; k++ )
	    {
		mesh->V[i][j][k] = mesh->V[i][j][k] - dt*((mesh->P[i-1][j][k-1]-mesh->P[i-1][j-1][k-1])/dy);
	    }
	}
    }

    /* CORRECT W */
#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,j,k) shared(mesh,dt,dx,dy,dz)
#endif
    for( i=1; i<mesh->nx; i++ )
    {
	for( j=1; j<mesh->ny; j++ )
	{
	    for( k=1;k<mesh->nz-2; k++ )
	    {
		mesh->W[i][j][k] = mesh->W[i][j][k] - dt*((mesh->P[i-1][j-1][k]-mesh->P[i-1][j-1][k-1])/dz);
	    }
	}
    }

    return 0;
}

//----------------------------------------BOUNDARY_COND
// BOUNDARY_COND
// COMPUTES BOUNDARY CONDITIONS AFTER STEP #1 AND STEP #2
// OF PROJECTION METHOD
// PRE: ALL VALUES ARE VALID
// POST: BOUNDARY CONDITIONS ARE SATISFIED
// RETURNS 0 ON SUCCESS, NONZERO OTHERWISE
//-----------------------------------------------------
int boundary_cond( double dt, 
		   double dx, 
		   double dy, 
		   double dz, 
		   NSMESH *mesh )
{
    /* VARS */
    int i = 0;
    int j = 0;
    int k = 0;
    /* ---- */
        
    /* WEST AND EAST FACE :: NO DELTA_X */
    #ifdef _OPENMP_
#pragma omp parallel for default(none) private(j,k) shared(mesh)
#endif
    for( j=0; j<mesh->ny; j++ )
    {
	for( k=0;k<mesh->nz; k++ )
	{
	    mesh->U[0][j][k] = 0; 
	    mesh->U[mesh->mx][j][k] = 0; 
	}
    }
    
#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,k) shared(mesh)
#endif
    /* NORTH AND SOUTH FACE :: NO DELTA_Y */
    for( i=0; i<mesh->nx; i++ )
    {
	for( k=0; k<mesh->nz; k++ )
	{
	    mesh->V[i][0][k] = 0;  
	    mesh->V[i][mesh->my][k] = 0; 
	}
    }

#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,j) shared(mesh)
#endif
    /* FRONT AND BACK FACE :: NO DELTA_Z */
    for( i=0; i<mesh->nx; i++ )
    {
	for( j=0; j<mesh->ny; j++ )
	{
	    mesh->W[i][j][0] = 0;  
	    mesh->W[i][j][mesh->mz] = 0; 
	}
    }

    /* SOLVE FINAL BOUNDARY CONDITIONS */
    /* -- U[:][1][:] = 2 * uS - U[:][2][:] */
    #ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,k) shared(mesh)
#endif
    // south 
    for( i=0; i<mesh->nx; i++ )
    {
        for( k=0; k<mesh->nz; k++ )
        {
            mesh->U[i][0][k] = 2 * 0. - mesh->U[i][1][k];
            mesh->W[i][0][k] = 2 * 0. - mesh->W[i][1][k];
        }
    }

#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,k) shared(mesh)
#endif
    // north
    /* -- U[:][ny][:] = 2 * uN - U[:][ny-1][:] */
    for( i=0; i<mesh->nx; i++ )
    {
        for( k=0; k<mesh->nz; k++ )
        {
	    mesh->U[i][mesh->ny-1][k] = 2 * 1. - mesh->U[i][mesh->ny-2][k];
	    mesh->W[i][mesh->ny-1][k] = 2 * 0. - mesh->W[i][mesh->ny-2][k];
        }
    }



#ifdef _OPENMP_
#pragma omp parallel for default(none) private(j,k) shared(mesh)
#endif
    /* -- V[1][:][:] = 2 * vW - V[2][:][:] */
    // west
    for( j=0; j<mesh->ny; j++ )
    {
        for( k=0; k<mesh->nz; k++ )
        {
            mesh->V[0][j][k] = 2 * 0. - mesh->V[1][j][k];
            mesh->W[0][j][k] = 2 * 0. - mesh->W[1][j][k];
        }
    }

    
#ifdef _OPENMP_
#pragma omp parallel for default(none) private(j,k) shared(mesh)
#endif
    /* -- V[nx][:][:] = 2 * vE - V[nx-1][:][:] */
    // east
    for( j=0; j<mesh->ny; j++ )
    {
        for( k=0; k<mesh->nz; k++ )
        {
	    mesh->V[mesh->nx-1][j][k] = 2 * 0 - mesh->V[mesh->nx-2][j][k];
	    mesh->W[mesh->nx-1][j][k] = 2 * 0 - mesh->W[mesh->nx-2][j][k];
        }
    }



#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,j) shared(mesh)
#endif
    /* -- W[:][:][1] = 2 * wF - W[:][:][2] */
    // back
    for( i=0; i<mesh->nx; i++ )
    {
        for( j=0; j<mesh->ny; j++ )
        {
            mesh->U[i][j][0] = 2 * 0 - mesh->U[i][j][1];
            mesh->V[i][j][0] = 2 * 0 - mesh->V[i][j][1];
        }
    }
    
#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,j) shared(mesh)
#endif
    /* -- W[:][:][nz] = 2 * wB - W[:][:][nz-1] */
    // front
    for( i=0; i<mesh->nx; i++ )
    {
        for( j=0; j<mesh->ny; j++ )
        {
	    mesh->U[i][j][mesh->nz-1] = 2 * 0 - mesh->U[i][j][mesh->nz-2];
	    mesh->V[i][j][mesh->nz-1] = 2 * 0 - mesh->V[i][j][mesh->nz-2];
        }
    }

    return 0;
}

//----------------------------------------VELSTAR
// VELSTAR
// COMPUTES INTERMEDIATE VELOCITY VALUES
// STEP #1 IN PROJECTION METHOD
// COMPUTES FORWARD LOOKING [N+1] VELOCITIES
// FOR U, V and W RESPECTIVELY
//
// MPP VERSION :: USES OPENMP TO PARALLELIZE 
//                VELOCITY UPDATE STEP
// 
// THE RESULTANT MESH [Ustar,Vstar,Wstar]
// IS TRANSFERED BACK TO THE ORIGINAL MESH [U,V,W]   
//
// PRE: MESH, DT, DX, DY, DZ, RE ARE VALID
// POST: STEP #1 IN PROJECTION METHOD AT TIME [t]
//       IS COMPLETE FOR ALL MESH POINTS
// RETURNS 0 ON SUCCESS, NONZERO OTHERWISE
//-----------------------------------------------
int velstar( double Re,
	     double dt, 
	     double dx, 
	     double dy, 
	     double dz, 
	     NSMESH *mesh )
{
    /* VARS */
    double dtdx  = dt/dx;
    double dtdy  = dt/dy;
    double dtdz  = dt/dz;
    double dtdxx = dt/(Re*dx*dx); 
    double dtdyy = dt/(Re*dy*dy); 
    double dtdzz = dt/(Re*dz*dz); 
    int i        = 0;
    int j        = 0;
    int k        = 0;
    /* ---- */
        
    /* -- U */
    for( i=1; i<mesh->mx; i++ )
    {
	for( j=1; j<mesh->ny-1; j++ )
	{
	    for( k=1; k<mesh->nz-1; k++ )
	    {
		mesh->Ustar[i][j][k] = mesh->U[i][j][k] 
		    - dtdx*(0.25*((mesh->U[i+1][j][k]+mesh->U[i][j][k])*(mesh->U[i+1][j][k]+mesh->U[i][j][k])) 
			    - 0.25*((mesh->U[i][j][k]+mesh->U[i-1][j][k])*(mesh->U[i][j][k]+mesh->U[i-1][j][k])))
		    - dtdy*(0.25*(mesh->U[i][j][k]+mesh->U[i][j+1][k])*(mesh->V[i][j][k]+mesh->V[i+1][j][k]) 
			    - 0.25*(mesh->U[i][j][k]+mesh->U[i][j-1][k])*(mesh->V[i][j-1][k]+mesh->V[i+1][j-1][k]))
		    - dtdz*(0.25*(mesh->U[i][j][k]+mesh->U[i][j][k+1])*(mesh->W[i][j][k]+mesh->W[i+1][j][k]) 
			    - 0.25*(mesh->U[i][j][k]+mesh->U[i][j][k-1])*(mesh->W[i][j][k-1]+mesh->W[i+1][j][k-1]))
		    + dtdxx*(mesh->U[i-1][j][k]-2*mesh->U[i][j][k]+mesh->U[i+1][j][k])
		    + dtdyy*(mesh->U[i][j-1][k]-2*mesh->U[i][j][k]+mesh->U[i][j+1][k])
		    + dtdzz*(mesh->U[i][j][k-1]-2*mesh->U[i][j][k]+mesh->U[i][j][k+1]);
	    }
	}
    }

    

    /* -- V */
    for( i=1; i<mesh->nx-1; i++ )
    {
	for( j=1; j<mesh->ny-2; j++ )
	{
	    for( k=1; k<mesh->nz-1; k++ )
	    {
		mesh->Vstar[i][j][k] = mesh->V[i][j][k]
                    - dtdy*(0.25*((mesh->V[i][j+1][k]+mesh->V[i][j][k])*(mesh->V[i][j+1][k]+mesh->V[i][j][k]))
                            -0.25*((mesh->V[i][j][k]+mesh->V[i][j-1][k])*(mesh->V[i][j][k]+mesh->V[i][j-1][k])))
                    - dtdx*(0.25*(mesh->V[i][j][k]+mesh->V[i][j+1][k])*(mesh->U[i][j][k]+mesh->U[i][j+1][k])
                            -0.25*(mesh->V[i][j][k]+mesh->V[i-1][j][k])*(mesh->U[i-1][j][k]+mesh->U[i-1][j-1][k]))
                    - dtdz*(0.25*(mesh->V[i][j][k]+mesh->V[i][j+1][k])*(mesh->W[i][j][k]+mesh->W[i][j][k+1])
                            -0.25*(mesh->V[i][j][k]+mesh->V[i][j][k-1])*(mesh->W[i][j][k-1]+mesh->W[i-1][j][k-1]))
                    + dtdxx*(mesh->V[i-1][j][k]-2*mesh->V[i][j][k]+mesh->V[i+1][j][k])
                    + dtdyy*(mesh->V[i][j-1][k]-2*mesh->V[i][j][k]+mesh->V[i][j+1][k])
                    + dtdzz*(mesh->V[i][j][k-1]-2*mesh->V[i][j][k]+mesh->V[i][j][k+1]);
	    }
	}
    }

    /* -- W */
    for( i=1; i<mesh->nx-1; i++ )
    {
	for( j=1; j<mesh->ny-1; j++ )
	{
	    for( k=1; k<mesh->nz-2; k++ )
	    {
		mesh->Wstar[i][j][k] = mesh->W[i][j][k]
                    - dtdz*(0.25*((mesh->W[i][j][k+1]+mesh->W[i][j][k])*(mesh->W[i][j][k+1]+mesh->W[i][j][k]))
                            -0.25*((mesh->W[i][j][k]+mesh->W[i][j][k-1])*(mesh->W[i][j][k]+mesh->W[i][j][k-1])))
                    - dtdx*(0.25*(mesh->W[i][j][k]+mesh->W[i][j][k+1])*(mesh->U[i][j][k]+mesh->U[i][j][k+1])
                            -0.25*(mesh->W[i][j][k]+mesh->W[i-1][j][k])*(mesh->U[i-1][j][k]+mesh->U[i-1][j][k-1]))
                    - dtdy*(0.25*(mesh->W[i][j][k]+mesh->W[i][j+1][k])*(mesh->V[i][j][k]+mesh->V[i][j][k+1])
                            -0.25*(mesh->W[i][j][k]+mesh->W[i][j-1][k])*(mesh->V[i][j-1][k]+mesh->V[i][j-1][k-1]))
                    + dtdxx*(mesh->W[i-1][j][k]-2*mesh->W[i][j][k]+mesh->W[i+1][j][k])
                    + dtdyy*(mesh->W[i][j-1][k]-2*mesh->W[i][j][k]+mesh->W[i][j+1][k])
                    + dtdzz*(mesh->W[i][j][k-1]-2*mesh->W[i][j][k]+mesh->W[i][j][k+1]);
	    }
	}
    }

    /* UPDATE ALL REAL VALUES */
#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,j,k) shared(mesh)
#endif
    for( i=0; i<mesh->nx; i++ )
    {
	for( j=0; j<mesh->ny; j++ )
	{
	    for( k=0; k<mesh->nz; k++ )
	    {
		mesh->U[i][j][k] = mesh->Ustar[i][j][k];
		mesh->V[i][j][k] = mesh->Vstar[i][j][k];
		mesh->W[i][j][k] = mesh->Wstar[i][j][k];
	    }		      
	}
    }
    
    return 0;
}

//----------------------------------------PRESSURE
// PRESSURE
// COMPUTES STEP #2 IN THE PROJECTION METHOD
// USING THE PRECOMPUTED MESH VELOCITIES; 
// COMPUTES PRESSURE AT EVERY CELL USING 
// SUCCESSIVE OVER RELAXATION METHOD [SOR]
// 
// PRE: DT, DX, DY, DZ and MESH ARE VALID
// POST: PRESSURE IS COMPUTED AT ALL CELLS
// RETURNS 0 ON SUCCESS, NONZERO OTHERWISE 
//------------------------------------------------
int pressure( double dt, 
	      double dx, 
	      double dy, 
	      double dz, 
	      NSMESH *mesh )
{
    /* VARS */
    double beta2y  = (dx*dx)/(dy*dy);
    double beta2z  = (dx*dx)/(dz*dz);
    double ap      = 2+(2*beta2y)+(2*beta2z);
    double ae      = beta2y;
    double aw      = beta2y;
    double an      = beta2z;
    double as      = beta2z;
    double af      = 1;
    double ab      = 1;
    double omega   = 1.6;
    double err_tol = 0.0001;     //-- ERROR TOLERANCE
    double err     = err_tol*2;  //-- ERROR
    double rhs     = 0;
    double tmpd    = 0;
    int iter       = 0;
    int i          = 0;
    int j          = 0;
    int k          = 0;
    int im1        = 0;
    int ip1        = 0;
    int jm1        = 0;
    int jp1        = 0;
    int km1        = 0;
    int kp1        = 0;
    /* ---- */

       
    /* BEGIN SOR SOLVER */
    while( (err>err_tol) && (iter<10000) )
    {
	for( i=0; i<mesh->mx; i++ )
	{
	    im1 = i-1;
	    ip1 = i+1;
	    if( i==0 )
	    {
		im1 = 0;
	    }
	    else if( i==mesh->mx-1 )
	    {
		ip1 = mesh->mx-1;
	    }
	    
	    for( j=0; j<mesh->my; j++ )
	    {
		jm1 = j-1;
		jp1 = j+1;

		if( j == 0 )
		{
		    jm1 = 0;
		}
		else if( j==mesh->my-1 )
		{
		    jp1 = mesh->my-1;
		}
		
		for( k=0; k<mesh->mz; k++ )
		{
		    km1 = k-1;
		    kp1 = k+1;

		    if( k==0 )
		    {
			km1 = 0;
		    }
		    else if( k==mesh->mz-1 )
		    {
			kp1==mesh->mz-1;
		    }
  
		    /* PRE :: CALCULATE RHS */
                    rhs =- dx*dx/dt * ( ((mesh->Ustar[i+1][j+1][k+1]-mesh->Ustar[i][j+1][k+1])/dx) 
                                        + ((mesh->Vstar[i+1][j+1][k+1]-mesh->Vstar[i+1][j][k+1])/dy) 
                                        + ((mesh->Wstar[i+1][j+1][k+1]-mesh->Wstar[i+1][j+1][k])/dz) );

		    /* STEP 1 */
                    mesh->Pnew[i][j][k] = ( (aw*mesh->Pnew[im1][j][k]) 
                                            + (ae*mesh->Pnew[ip1][j][k]) 
                                            + (as*mesh->Pnew[i][jm1][k]) 
                                            + (an*mesh->Pnew[i][jp1][k]) 
                                            + (af*mesh->Pnew[i][j][km1]) 
                                            + (ab*mesh->Pnew[i][j][kp1]) 
                                            + rhs )/ap;
		    
		    /* STEP 2 */
		    mesh->Pnew[i][j][k] = omega*mesh->Pnew[i][j][k]+(1-omega)*mesh->P[i][j][k];
		}
	    }
	}
	
	/* FIND MAXIMUM ERROR */
#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,j,k,tmpd) shared(mesh,err)
#endif
	for( i=0; i<mesh->mx; i++ )
	{
	    for( j=0; j<mesh->my; j++ )
	    {
		for( k=0; k<mesh->mz; k++ )
		{
		    tmpd = abs(mesh->Pnew[i][j][k] - mesh->P[i][j][k]);
		    if( err < tmpd )
		    {
			err = tmpd;
		    }
		}
	    }
	}
    
#ifdef _OPENMP_
#pragma omp parallel for default(none) private(i,j,k) shared(mesh)
#endif
	/* COPY PNEW TO P */
	for( i=0; i<mesh->mx; i++ )
	{
	    for( j=0; j<mesh->my; j++ )
	    {
		for( k=0; k<mesh->mz; k++ )
		{
		    mesh->P[i][j][k] = mesh->Pnew[i][j][k];
		}
	    }
	}
	
	iter++;
    }

    return 0;
}

//----------------------------------------STREAM
// STREAM
//----------------------------------------------
int stream( int nx, 
	    int ny, 
	    int nz, 
	    double dx, 
	    double dy, 
	    double dz, 
	    NSMESH *mesh )
{
    return 0;
}

/* EOF */
